import numpy as np
import matplotlib.pyplot as plt
from BernoulliBandit import Solver, bandit_10_arm, plot_results

# 随时间衰减的epsilon-贪婪算法
class DecayingEpsilonGreedy(Solver):
    def __init__(self, bandit, tr=0.01, init_prob=1.0):
        super(DecayingEpsilonGreedy, self).__init__(bandit)
        self.estimates = np.array([init_prob] * self.bandit.K)  # 初始化拉动每根拉杆的期望奖励估值
        self.steps = 0

    def run_one_step(self):
        self.steps += 1
        if np.random.random() < 1/self.steps:
            k = np.random.randint(0, self.bandit.K)  # 随机选择一根拉杆
        else:
            k = np.argmax(self.estimates)  # 选择期望奖励估值最大的拉杆
        r = self.bandit.step(k)  # 得到本次动作的奖励
        self.estimates[k] += 1. / (self.counts[k] + 1) * (r - self.estimates[k])
        return k

def run_single():
    np.random.seed(1)
    epsilon_greedy_solver = DecayingEpsilonGreedy(bandit_10_arm, tr=0.01)
    times = 5000
    epsilon_greedy_solver.run(times)
    print('epsilon-贪婪算法的累积懊悔为：', epsilon_greedy_solver.regret)
    plot_results([epsilon_greedy_solver], ["DecayingEpsilonGreedy"])


if __name__ == "__main__":
    run_single()

